# -*- coding: utf-8 -*-

import socket
import re
import json
import requests
from flask import Flask
import dns.resolver
            
def whois_request(domain: str, server: str, port=43, timeout=5) -> str:
    """
    发送http请求，获取信息
    :param domain:
    :param server:
    :param port:
    :return:
    """
    # 创建连接
    sock = socket.create_connection((server, port))
    sock.settimeout(timeout)

    # 发送请求
    sock.send(("%s\r\n" % domain).encode("utf-8"))

    # 接收数据
    buff = bytes()
    while True:
        data = sock.recv(1024)
        if len(data) == 0:
            break
        buff += data

    # 关闭链接
    sock.close()

    return buff.decode("utf-8")

def getValue(response : str,key: str):
    matchObj  = re.search( r'{}: (.*)'.format(key), response, re.M|re.I)
    if matchObj :
        match = matchObj.group()
        k , v = match.split(':')
        return v.strip()
    return None

def whois(domain):
    # 根服务器地址
    root_server = 'whois.iana.org'
    # 查询域名信息
    res = whois_request(domain, root_server)
    refer = getValue(res,'refer')
    print('refer:', refer)
    if not refer is None:
        true_res = whois_request(domain, refer)
        # print(true_res)
        return true_res
    return None

def resolveResponse(response : str):
    keys = ['Domain Name',
        'Registry Domain ID',
        'Registrar WHOIS Server',
        'Registrar URL',
        'Updated Date',
        'Creation Date',
        'Registry Expiry Date',
        'Registrar',
        'Registrar IANA ID',
        'Registrar Abuse Contact Email',
        'Registrar Abuse Contact Phone',
        'Domain Status',
        'Name Server']
    ret = {}
    for key in keys:
        ret[key] = ''
    response = response.split('\r\n')
    for row in response:
        row_list = row.split(':',1)
        if len(row_list) == 2:
            key = row_list[0].strip()
            if key in keys:
                if len(ret[key]) == 0:
                    ret[key] = str(row_list[1]).strip()
                else:
                    ret[key] = ret[key] + ", " + str(row_list[1]).strip()
    return ret

def queryDns(domain):
    result_a = dns.resolver.resolve(domain, 'A')  #查询记录为A记录
    ret = []
    for i in result_a.response.answer:
        for j in i.items:
            if j.rdtype == 1: #加判断，不然会出现AttributeError: 'CNAME' object has no attribute 'address'
                print(j.address)
                ret.append(j.address)
    return ret
            
app = Flask(__name__)


@app.route("/")
@app.route("/whois/<domain>/")
def r_whois( domain ):
    info = whois(domain)
    if info is None:
        return json.dumps({'code':40400,'message': domain + '查无数据'}), 200, {"Content-Type":"application/json;charset=UTF-8"}
    else:
        whois_info = resolveResponse(info)
        return json.dumps({'code':20000, 'message': '成功', 'data':whois_info}), 200, {"Content-Type":"application/json;charset=UTF-8"}
        
@app.route("/dns/<domain>/")
def r_dns( domain ):
    info = queryDns(domain)
    if len(info) == 0:
        return json.dumps({'code':40400,'message': domain + '查无数据'}), 200, {"Content-Type":"application/json;charset=UTF-8"}
    else:
        return json.dumps({'code':20000, 'message': '成功', 'data':info}), 200, {"Content-Type":"application/json;charset=UTF-8"}        

@app.route("/host/<ip>/")
def r_shodan( ip ):
    # 调用shodan接口查询IP地址的信息
    api_key = "ivpQMs45TYimeuWqW97anU4BMYahsCsl";
    url = "https://api.shodan.io/shodan/host/{}?key={}"
    resp = requests.get(url.format(ip,api_key),timeout = 10)
    ret = {}
    ret["code"] = 20000
    ret["data"] = json.loads(resp.text)
    return json.dumps(ret), 200, {"Content-Type":"application/json;charset=UTF-8"}
    
@app.route("/user/login", methods=['POST'])
def r_user_login():
    ret = {
        "code": 20000,
        "data": 'admin-token'
    }
    return json.dumps(ret), 200, {"Content-Type":"application/json;charset=UTF-8"} 
    
@app.route("/user/info")
def r_user_info():
    ret = {
        "code": 20000,
        "data": {
            "roles": ['admin'],
            "introduction": 'I am a super administrator',
            "avatar": 'https://wpimg.wallstcn.com/f778738c-e4f8-4870-b634-56703b4acafe.gif',
            "name"  : 'Super Admin'
        }
    }
    return json.dumps(ret), 200, {"Content-Type":"application/json;charset=UTF-8"} 
    
    
if __name__ == '__main__':
    info = whois('baidu.com')
    whois_info = resolveResponse(info)
    print(whois_info)
    iplist = queryDns('baidu.com')
    print(iplist)